package cim144

import chisel3._ // VecInit
import chisel3.util._ // MuxCase
import freechips.rocketchip.tile._
import freechips.rocketchip.config._
import freechips.rocketchip.diplomacy._ // LazyModule
import freechips.rocketchip.rocket._ // 

class CIM144(opcode: OpcodeSet)(implicit p: Parameters) extends LazyRoCC(opcode) with SystemConfig {
    override lazy val module = new LazyRoCCModuleImp(this) {
        //******1.FSM*****************
        val idle::push ::save::mvm::config::Nil = Enum(5)
        val state = RegInit(idle)
        val isidle = state === idle
        val ispush = state === push
        val ismvm  = state === mvm
        val issave = state === save
        val isconfig = state === config
        //*****1.1 decode*************
        val rv32_func7_0010001 = io.cmd.bits.inst.funct === "b0010001".U
        val rv32_func7_0010010 = io.cmd.bits.inst.funct === "b0010010".U
        val rv32_func7_0000001 = io.cmd.bits.inst.funct === "b0000001".U
        val rv32_func7_0000010 = io.cmd.bits.inst.funct === "b0000010".U
        val custom_push  = rv32_func7_0010001 & isidle & io.cmd.valid & io.cmd.ready
        val custom_save  = rv32_func7_0010010 & isidle & io.cmd.valid & io.cmd.ready
        val custom_mvm   = rv32_func7_0000001 & isidle & io.cmd.valid & io.cmd.ready
        val custom_config= rv32_func7_0000010 & isidle & io.cmd.valid & io.cmd.ready
        //*********** 1.2 finite state machine***********
        val push_done = Wire(Bool());val save_done = Wire(Bool()); val mvm_done = Wire(Bool());
        switch(state){
            is(idle){ state :=  Mux(custom_push,push,
                                Mux(custom_save,save,
                                Mux(custom_mvm,mvm,
                                Mux(custom_config,config,idle))))}
            is(push){ state :=  Mux(push_done,idle,push)}
            is(save){ state :=  Mux(save_done,idle,save)}
            is(mvm){  state :=  Mux(mvm_done ,idle,mvm)}
            is(config){state:= idle}
        }
        // ********* 1.3 some output signals **********
        //        val config_done = isconfig
        io.busy := Mux(isidle,io.cmd.valid,true.B) // io.busy = !isidle
        io.cmd.ready := isidle
        //********** 2 Datapath ******************
        //**********2.0 config parameter
        //rs1: *******23 |22********16|15**********8|7*********0|
        //     bit_random|**wise_num**|**out_shift**|**adc_range|
        val con_para = Wire(new config_para)
        con_para.adc_range  := RegEnable(io.cmd.bits.rs1(7 , 0),0.U,custom_config)
        con_para.out_shift  := RegEnable(io.cmd.bits.rs1(15, 8),0.U,custom_config)
        con_para.wise_num   := RegEnable(io.cmd.bits.rs1(22,16),0.U,custom_config)
        con_para.bit_random := RegEnable(io.cmd.bits.rs1(23),false.B,custom_config)
        // *********2.1 push and save access time generate************
        // rs2: *********|21********11|10********0|
        //      *********|start_index*|*push_num**|
        //      start_index and push_num should be Multiples of 8
        // rs1: |63****************memory addr***0|
        val buf_end_index       = RegInit(0.U(log2Ceil(PUSH_MAX_SIZE+1).W))
        val push_save_buf_cnt   = RegInit(0.U(log2Ceil(PUSH_MAX_SIZE+1).W))   // 7bit
        //val save_buf_cnt  = RegInit(0.U(log2Ceil(SAVE_MAX_SIZE+1).W))
        val baseAddr      = RegInit(0.U(CIM_XLEN.W))
        // buf is 64bits, data is 8bit,
        //buf_end_index := RegEnable(io.cmd.bits.rs2(21,14) + io.cmd.bits.rs2(10,3) -1.U , custom_push|custom_save)
        // push_buf_cnt and save_buf_cnt can reuse?
        when(custom_push|custom_save){
            buf_end_index := io.cmd.bits.rs2(21,14)  +  io.cmd.bits.rs2(10,3) -1.U
            push_save_buf_cnt  := io.cmd.bits.rs2(21,14)
            baseAddr      := io.cmd.bits.rs1
        }
        push_done := (push_save_buf_cnt === buf_end_index) & ispush & io.mem.resp.valid
        // 1.IMPORTANT only when mem ready to accept the data, can "save_done" finish
        save_done := (push_save_buf_cnt === buf_end_index) & issave & io.mem.req.ready
        //**********2.2 push buffer generate*********
        val push_buf    = RegInit(VecInit(Seq.fill(PUSH_MAX_SIZE)(0.U(CIM_XLEN.W))))  //72 x 64 bit
        //**********2.3 bitwise module *************
        val push_buf_8bit = push_buf.asTypeOf(Vec(ROW_NUM,SInt(8.W)))
        val push_buf_8bit_usign = VecInit(Seq.fill(ROW_NUM)(0.U(7.W)))
        val counter_576 = RegInit(VecInit(Seq.fill(ROW_NUM)(0.U(7.W))))
        val compare_576 = VecInit(Seq.fill(ROW_NUM)(false.B))
        val push_buf_sel = VecInit(Seq.fill(ROW_NUM)(0.U(2.W)))
        val cim_done = WireInit(false.B)
        for(i<-0 until ROW_NUM){
            push_buf_8bit_usign(i) := Mux(push_buf_8bit(i)(7),-push_buf_8bit(i),push_buf_8bit(i)).asUInt
            compare_576(i) := push_buf_8bit_usign(i) > counter_576(i)
            push_buf_sel(i):= Cat(compare_576(i)&push_buf_8bit(i)(7),compare_576(i)& !push_buf_8bit(i)(7))
            when(custom_mvm){
                counter_576(i) := Mux(con_para.bit_random,(i%128).U,0.U)
            }.elsewhen(cim_done){
                counter_576(i) := counter_576(i) + 1.U
            }
        }
        val push_1152 = push_buf_sel.asUInt()
        // **********2.4 mvm process *****************
        // **********2.4.1 cim cell ******************
        val bitwise_end = (counter_576(0)===con_para.wise_num) & ismvm
        val mvm_start_tmp = RegNext(Mux(custom_mvm,true.B,Mux(cim_done&ismvm,true.B,false.B)))
        val mvm_start     = RegNext(Mux(mvm_done,false.B,mvm_start_tmp))
        val rcbd = Wire(new mvm_IO())
        //rs1 |31******************16|15***************0|
        //    |*******row_num********|******col_num*****|
        //rs2 |31*********22|21*********11|10**********0|
        //    |*array index*|**start_row**|**start_col**|
        rcbd.row_begin := RegEnable(io.cmd.bits.rs2(21,11),0.U,custom_mvm)
        rcbd.row_end   := RegEnable(io.cmd.bits.rs1(31,16)+io.cmd.bits.rs2(21,11) - 1.U,0.U,custom_mvm)
        rcbd.col_begin := RegEnable(io.cmd.bits.rs2(10, 0),0.U,custom_mvm)
        rcbd.col_end   := RegEnable(io.cmd.bits.rs1(15, 0)+io.cmd.bits.rs2(10, 0) - 1.U,0.U,custom_mvm)
        val cim_out   = cim_mvm_8bit_shift(mvm_start,rcbd,push_1152,con_para.adc_range)
        cim_done  := RegNext(cim_out._1)
        val cim_result = cim_out._2.asTypeOf(Vec(COL_NUM,SInt(8.W)))
        //**********2.4.2 acc cim_result ***************
        val mvm_tmp         = RegInit(VecInit(Seq.fill(COL_NUM)(0.S(16.W))))
        val mvm_tmp_shift   = VecInit(Seq.fill(COL_NUM)(0.S(16.W)))
        val mvm_result      = VecInit(Seq.fill(COL_NUM)(0.S(8.W)))
        for(i<-0 until COL_NUM) {
            when(custom_mvm) {
                mvm_tmp(i) := 0.S
            }.elsewhen(cim_done) {
                mvm_tmp(i) := mvm_tmp(i) + cim_result(i)
            }
            mvm_tmp_shift(i) := mvm_tmp(i) >> con_para.out_shift
            mvm_result(i) := Mux(mvm_tmp_shift(i) >= 127.S, 127.S, Mux(mvm_tmp_shift(i) <= -127.S, -127.S, mvm_tmp_shift(i)))
        }
        // 2.IMPORTANT mem signal!
        //mvm_done := ismvm & cim_done & bitwise_end
        mvm_done := ismvm & bitwise_end
        //********** 2.5 save wire generate*********
        //val save_wire   = VecInit(Seq.fill(SAVE_MAX_SIZE)(0.U(CIM_XLEN.W)))
        val save_16x64 = mvm_result.asTypeOf(Vec(SAVE_MAX_SIZE,UInt(CIM_XLEN.W)))
        // ***********3.memory access, push and save *******
        //3.1 ********* mem io setting***********
        val memReq = io.mem.req
        val memResp = io.mem.resp
        memReq.bits.cmd := Mux(issave, M_XWR, M_XRD)
        memReq.bits.addr := baseAddr
        //memReq.bits.size := log2Ceil(64).U
        //memReq.bits.size := 6.U
        //3. IMPORTANT! 0.U-> 1Byte, 1.U->2Byte , 2.U -> 4Byte, 3.U -> 8Byte
        memReq.bits.size := 3.U

        val mem_idle :: mem_op :: Nil= Enum(2)
        val mem_state = RegInit(mem_idle)
        memReq.valid := (ispush || issave) && (mem_state===mem_idle) && !(push_done)
        
        //4. IMPORTANT! mem_state should back to "mem_idle" after "ROCC_save"
        when(isidle){
            mem_state := mem_idle
        }.elsewhen(memReq.fire()){
            mem_state := mem_op
        }
        memReq.bits.data := save_16x64(push_save_buf_cnt)
        // 3.2*************cache read**********
        when(mem_state===mem_op && ispush && memResp.valid){
            baseAddr                := baseAddr + 8.U
            push_save_buf_cnt            := push_save_buf_cnt + 1.U
            push_buf(push_save_buf_cnt)  := memResp.bits.data
            mem_state := mem_idle
        }
        // 3.3************* cache write**********
        when(mem_state===mem_op && issave  ){
            baseAddr := baseAddr + 8.U
            push_save_buf_cnt := push_save_buf_cnt + 1.U
            mem_state := mem_idle
        }
    }

}

class CIM144_rocc extends Config((site,here,up) => {
    case BuildRoCC => Seq(
        (p:Parameters) => {
            //val regBufferNum = 128 // RoCC buffer numer
            val cim144 = LazyModule(new CIM144(OpcodeSet.custom0)(p))
            cim144
        }
    )
})

